#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""

@Time   :2025/8/10 上午11:05
@Author :zengjiahao1989@gmail.com
@File   :3. RunnablePassthrough简化invoke.py
"""

import os

import dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()


def retrieval(query: str) -> str:
    """一个模拟的检索器函数"""
    print("正在检索query")
    return "我是曾嘉昊"


# 1.编排prompt
prompt = ChatPromptTemplate.from_template("""根据用户的问题回答，可以参考上下文进行生成
<content>
{context}
</content>


用户的提问就是:{query}""")
# 2.构建大语言模型
llm = ChatOpenAI(
    model_name="kimi-k2-0711-preview",
    openai_api_key=os.getenv("OPENAI_API_KEY"),
    openai_api_base=os.getenv("OPENAI_API_BASE"),
)
# 3.创建链 RunnablePassthrough获取传入参数，字符串返回字符串
# chain = {"query": RunnablePassthrough()} | prompt | llm | StrOutputParser()
# 传入dist返回dist，并且可以通过assign方法添加一个dist的参数
chain = RunnablePassthrough.assign(context=lambda x: retrieval(x['query'])) | prompt | llm | StrOutputParser()
# 4.调用链并获取结果
# 传入参数为字符串，就返回字符串
# content = chain.invoke("你好，你是")

# 传入参数是dist，就返回dist
content = chain.invoke({"query": "你好，我是谁"})
print(content)
